import os
from io import BytesIO

import fitz
import numpy as np
from PIL import Image
from fastapi import UploadFile

from img.ocr_img import ImgParser
from output.tools import get_data_text
from tbpu import get_parser
from tbpu.parser_tools.paragraph_parse import word_separator
from server import FileParser, DEFAULT_TBPU, UPLOAD_DIR
from util.pdf import parse_pdf


class PdfParser(FileParser):
    extensions = ['pdf']

    async def parse_upload_file_to_text(self, file: UploadFile, tbpu=DEFAULT_TBPU):
        local_file = os.path.join(UPLOAD_DIR, file.filename)
        with open(local_file, "wb") as f:
            content = await file.read()
            f.write(content)
        resp = self.parse_file_to_text(local_file, tbpu)
        os.remove(local_file)
        return resp

    def parse_file_to_text(self, filename, tbpu=DEFAULT_TBPU):
        return self.parse_pdf_text(filename, tbpu)

    def parse_pdf_text(self, pdf_path, tbpu):
        img_parser = ImgParser()
        page_arr = parse_pdf(pdf_path, tbpu, img_parser.parse_img_nd)
        paser = get_parser(tbpu)
        resp = ""
        # 计算平均置信度
        score, num = 0, 0
        for page in page_arr:
            for r in page:
                score += r["score"]
                num += 1
            if num > 0:
                score /= num

            # 执行 tbpu
            if paser is not None and len(page) > 0:
                page = paser.run(page)
            resp += "\n" + get_data_text(page)

        return {
            'score': score,
            'data': resp
        }

    # def parse_pdf(self, pdf_path, tbpu, page_start=None, page_end=None):
    #     try:
    #         doc = fitz.open(pdf_path)
    #     except Exception as e:
    #         print(f"[Error] fitz.open error: {pdf_path}, Message: {e}")
    #         return []
    #     # if doc.is_encrypted and not doc.authenticate(password):
    #     #     if password:
    #     #         msg = f"[Error] Incorrect password. 文档已加密，密码错误。 [{password}]"
    #     #     else:
    #     #         msg = "[Error] Doc encrypted. 文档已加密，请提供密码。"
    #     #     return msg
    #     page_count = doc.page_count
    #     if page_start is not None and page_end is not None:
    #         if page_start < 0:
    #             page_start += page_count + 1
    #         if page_end < 0:
    #             page_end += page_count + 1
    #         if page_start < 1:
    #             return f"[Error] 起始页码不能小于1"
    #         if page_end > page_count:
    #             return f"[Error] 结束页码不能大于页数"
    #         if page_start > page_end:
    #             return f"[Error] 页码范围错误"
    #         page_list = list(range(page_start - 1, page_end))
    #     else:
    #         page_list = list(range(0, page_count))
    #
    #     tb_arr = []
    #     for page_num in page_list:
    #         page = doc[page_num]
    #         resp = self.parse_single_page(page, page_num, tbpu)
    #
    #         tb_arr.append(resp)
    #
    #     return tb_arr
    #
    # def parse_single_page(self, page, page_num, tbpu):
    #     imgs = []  # 待OCR的图片列表
    #     tbs = []  # text box 文本块列表
    #     page_rotation = page.rotation
    #     # 获取元素 https://pymupdf.readthedocs.io/en/latest/_images/img-textpage.png
    #     # 确保越界图像能被采集 https://github.com/pymupdf/PyMuPDF/issues/3171
    #     p = page.get_text("dict", clip=fitz.INFINITE_RECT())
    #     for t in p["blocks"]:  # 遍历区块（段落）
    #         # 图片
    #         if t["type"] == 1:  # 图片内容
    #             img_bytes = t["image"]  # 图片字节
    #             bbox = t["bbox"]  # 图片包围盒
    #             # 图片视觉大小、原始大小、缩放比例
    #             w1, h1 = bbox[2] - bbox[0], bbox[3] - bbox[1]
    #             w2, h2 = t["width"], t["height"]
    #             # 特殊情况：图片宽高为0
    #             if w2 <= 0 or h2 <= 0:
    #                 continue
    #             # 单独计算宽高的缩放比例
    #             scale_w = w1 / w2
    #             scale_h = h1 / h2
    #             nd_arr = np.array(None)
    #             try:
    #                 with Image.open(BytesIO(img_bytes)) as pimg:
    #                     # 记录原图格式
    #                     format = pimg.format
    #                     if not format:
    #                         format = "PNG"
    #                     if page_rotation != 0:
    #                         # 如果页面有旋转，逆向旋转图片字节
    #                         print(f"P{page_num} - 旋转 {page_rotation} °")
    #                         # PDF的旋转是顺时针，需要逆时针旋转图片
    #                         pimg = pimg.rotate(-page_rotation, expand=True)
    #                         # 将旋转后的图片转回bytes
    #                         buffered = BytesIO()
    #                         pimg.save(buffered, format=format)
    #                         img_bytes = buffered.getvalue()
    #                     nd_arr = np.array(pimg)
    #             except Exception:
    #                 print("转换文档图片异常。")
    #             # 记录图片
    #             imgs.append(
    #                 {
    #                     "nd_arr": nd_arr,
    #                     "bytes": img_bytes,
    #                     "xy": (bbox[0], bbox[1]),
    #                     "scale_w": scale_w,
    #                     "scale_h": scale_h,
    #                     "end": "\n"
    #                 }
    #             )
    #         # 文本
    #         elif t["type"] == 0:
    #             l = len(t["lines"]) - 1
    #             for index, line in enumerate(t["lines"]):  # 遍历每一行
    #                 text = ""
    #                 for span in line["spans"]:  # 遍历每一文本块
    #                     text += span["text"]
    #                 if text:
    #                     b = line["bbox"]
    #                     tb = {
    #                         "box": [
    #                             [b[0], b[1]],
    #                             [b[2], b[1]],
    #                             [b[2], b[3]],
    #                             [b[0], b[3]],
    #                         ],
    #                         "text": text,
    #                         "score": 1,
    #                         "end": "\n" if index == l else "",  # 结尾符
    #                         "from": "text",  # 来源：直接提取文本
    #                     }
    #                     tbs.append(tb)
    #     # 补充结尾符
    #     for i1 in range(len(tbs) - 1):
    #         if tbs[i1]["end"]:  # 跳过已有结尾符的
    #             continue
    #         i2 = i1 + 1
    #         sep = word_separator(tbs[i1]["text"][-1], tbs[i2]["text"][0])
    #         tbs[i1]["end"] = sep
    #
    #     img_parser = ImgParser()
    #
    #     if imgs:
    #         for img in imgs:
    #             res = img_parser.parse_img_nd(img['nd_arr'])
    #             if res is not None:
    #                 x, y = img["xy"]
    #                 scale_w = img["scale_w"]
    #                 scale_h = img["scale_h"]
    #                 for r in res:
    #                     # 将所有文本块的坐标，从图片相对坐标系，转为页面绝对坐标系
    #                     for bi in range(4):
    #                         r["box"][bi][0] = r["box"][bi][0] * scale_w + x
    #                         r["box"][bi][1] = r["box"][bi][1] * scale_h + y
    #                     r["from"] = "img"  # 来源：OCR
    #                     tbs.append(r)
    #     parser = get_parser(tbpu)
    #     if parser is not None and len(tbs) > 0:
    #         tbs = parser.run(tbs)
    #     return tbs
